#!/bin/bash

# --- 配置区 ---
# 设置你的 .pth 文件所在的目录路径
TARGET_DIR="/home/zhangyichi/projects/DeepDR_Plus/logs/bs16_epoch150_lr0.0001/TrainerDR_20250409_152323"  # <<< --- 修改这里：替换成你的实际目录路径

# 设置阈值 Epoch (文件名中的数字小于这个值的文件将被删除)
THRESHOLD_EPOCH=149                 # <<< --- 修改这里：比如你想删除 epoch 小于 50 的文件

# 设置为 "true" 以进行预览 (只打印不删除)，设置为 "false" 以执行实际删除操作
# *** 务必先用 true 运行一次，确认无误后再改成 false ***
DRY_RUN="false"                       # <<< --- 安全起见，先保持 true
if [ "$DRY_RUN" != "false" ]; then
    DRY_RUN="true" # Ensure it's either true or false, default to true (safe mode)
fi
# --- 配置结束 ---

# 检查目录是否存在
if [ ! -d "$TARGET_DIR" ]; then
  echo "错误：目录 '$TARGET_DIR' 不存在。"
  exit 1
fi

echo "将在目录 '$TARGET_DIR' 中操作..."
if [ "$DRY_RUN" = "true" ]; then
  echo "*** 当前为预览模式 (DRY RUN)，不会实际删除文件 ***"
  echo "将查找 Epoch < $THRESHOLD_EPOCH 的文件进行模拟删除："
else
  echo "*** 警告：当前为实际删除模式 ***"
  echo "将删除 Epoch < $THRESHOLD_EPOCH 的文件："
fi
echo "--------------------------------------------------"

# 遍历目录下的所有 .pth 文件
shopt -s nullglob # 如果没有匹配的文件，让循环不执行而不是报错
for file_path in "$TARGET_DIR"/*.pth; do
  # 提取文件名（不含路径）
  filename=$(basename "$file_path")

  # 移除 .pth 后缀
  base_name="${filename%.pth}"
  epoch_str="" # 初始化 Epoch 字符串

  # 使用正则表达式尝试提取 Epoch 数字
  # 这个正则匹配两种模式:
  # 1. 文件名就是纯数字 (e.g., "50")
  # 2. 文件名以 "_数字" 结尾 (e.g., "model_120", "checkpoint_30")
  # 它会捕获最后一部分连续的数字
  if [[ "$base_name" =~ ^(.*_)?([0-9]+)$ ]]; then
      # ^ : 匹配字符串开头
      # (.*_)? : 可选地匹配任意字符后面跟着一个下划线 (非捕获或捕获都行，这里捕获了前缀部分到 BASH_REMATCH[1])
      # ([0-9]+) : 捕获一个或多个数字 (这是我们想要的 Epoch，在 BASH_REMATCH[2])
      # $ : 匹配字符串结尾

      epoch_str="${BASH_REMATCH[2]}" # 提取捕获到的数字部分

      # 确保提取到的是非空数字字符串 (理论上正则保证了，但加个检查更稳妥)
      if [[ -n "$epoch_str" ]]; then
          # 进行数值比较
          if [ "$epoch_str" -lt "$THRESHOLD_EPOCH" ]; then
              # Epoch 小于阈值
              if [ "$DRY_RUN" = "true" ]; then
                  echo "[预览] 会删除: $filename (Epoch: $epoch_str)"
              else
                  echo "正在删除: $filename (Epoch: $epoch_str)"
                  rm "$file_path"
              fi
          # else
              # 可选：打印那些 Epoch 大于等于阈值的文件
              # echo "[保留] Epoch >= $THRESHOLD_EPOCH: $filename (Epoch: $epoch_str)"
          fi
      else
          # 正则匹配了但数字部分为空，理论上不应发生
           echo "[跳过] 无法从 '$filename' 中解析有效 Epoch 数字 (匹配但数字为空)。"
      fi
  else
      # 文件名格式不符合预期的两种模式 (纯数字.pth 或 前缀_数字.pth)
      echo "[跳过] 文件名格式无法解析 Epoch: $filename"
  fi
done
shopt -u nullglob # 关闭 nullglob

echo "--------------------------------------------------"
if [ "$DRY_RUN" = "true" ]; then
  echo "*** 预览模式结束。检查以上列表，如果无误，请编辑脚本将 DRY_RUN 设置为 false 再运行以实际删除。 ***"
else
  echo "*** 删除操作完成。 ***"
fi